from dataclasses import dataclass
from typing import Dict, Optional, List
from .ast_nodes import *

@dataclass
class Symbol:
    name: str
    type: TypeNode
    is_mutable: bool = False

class SymbolTable:
    def __init__(self, parent: Optional['SymbolTable'] = None):
        self.symbols: Dict[str, Symbol] = {}
        self.parent = parent

    def define(self, name: str, type: TypeNode, is_mutable: bool = False):
        self.symbols[name] = Symbol(name, type, is_mutable)

    def lookup(self, name: str) -> Optional[Symbol]:
        if name in self.symbols:
            return self.symbols[name]
        if self.parent:
            return self.parent.lookup(name)
        return None

class TypeChecker:
    def __init__(self):
        self.global_scope = SymbolTable()
        self.types: Dict[str, RecordType] = {} # Store user defined types

        # 初始化内置类型和函数
        # 现在，只基本类型
        self.int_type = TypeNode("数", [])
        self.float_type = TypeNode("徽数", [])
        self.bool_type = TypeNode("两仪", [])
        self.string_type = TypeNode("词", [])
        self.void_type = TypeNode("空", [])
        
        # Add printf to global scope
        # printf : (词) -> 数
        printf_type = FunctionType("Function", [], [self.string_type], self.int_type)
        self.global_scope.define("printf", printf_type)
        
        # Add 书 (print) - 多态打印函数
        # 现在，我们将其类型定义为接受任何单个参数（使用 void 作为占位符表示 "any"）
        # 书 : (T) -> 空
        print_type = FunctionType("Function", [], [self.void_type], self.void_type)
        self.global_scope.define("书", print_type)

    def check_program(self, program: Program):
        for defn in program.definitions:
            self.check_definition(defn, self.global_scope)

    def check_definition(self, node: ASTNode, scope: SymbolTable):
        if isinstance(node, RecordDef):
            # 定义记录类型
            fields = {}
            for field in node.fields:
                # 确保字段类型有效（这里简化，假设所有类型名都有效或已定义）
                # 实际应检查 field.type_annotation 是否存在于 self.types 或内置类型中
                fields[field.name] = field.type_annotation
            
            record_type = RecordType(node.name, [], fields)
            self.types[node.name] = record_type
            # 也可以将其添加到作用域中作为类型符号，如果我们需要在运行时引用类型
            
        elif isinstance(node, VariableDef):
            # 检查变量类型
            value_type = self.check_expression(node.value, scope)
            print(f"Variable '{node.name}' type: {value_type}")
            
            # 检查注释如果存在
            if node.type_annotation:
                # 简单名称检查
                if node.type_annotation.name != value_type.name:
                    raise TypeError(f"Type mismatch for variable '{node.name}': expected {node.type_annotation.name}, got {value_type.name}")
            
            # 定义在作用域中
            scope.define(node.name, value_type, node.is_mutable)
            
        elif isinstance(node, FunctionDef):
            print(f"DEBUG sema: Checking function '{node.name}'")
            # 创建函数作用域
            func_scope = SymbolTable(scope)
            # 检查参数类型
            param_types = []
            for param in node.params:
                param_type = self.void_type
                if param.type_annotation:
                    # 假设 map_type_node 是 check_expression 的变体，或者直接用 check_expression 如果它是 Expr
                    # 但 type_annotation 是 TypeNode 还是 Expr?
                    # parser 返回 TypeNode (ASTNode)? No, parser returns TypeNode object directly?
                    # Let's check parser.py later. Assuming param.type_annotation is TypeNode.
                    # Wait, sema uses TypeNode class.
                    # If param.type_annotation is already TypeNode, use it.
                    param_type = param.type_annotation
                else:
                    param_type = self.int_type # Default to Int
                
                param_types.append(param_type)
                func_scope.define(param.name, param_type)
            
            # 如果有返回类型，先定义函数以支持递归
            if node.return_type:
                # 预定义
                func_type = FunctionType("Function", [], param_types, node.return_type)
                scope.define(node.name, func_type)

            # 检查函数体
            body_type = self.check_block(node.body, func_scope)
            print(f"DEBUG sema: Function '{node.name}' inferred body type: {body_type.name}")
            
            # 检查返回类型一致性
            if node.return_type:
                if node.return_type.name != body_type.name:
                    raise TypeError(f"函数 '{node.name}' 的返回类型不匹配：期望 {node.return_type.name}，但得到 {body_type.name}")
            else:
                # 如果没有注解，现在定义（不支持递归推断）
                func_type = FunctionType("Function", [], param_types, body_type)
                scope.define(node.name, func_type)

    def check_expression(self, node: Expr, scope: SymbolTable) -> TypeNode:
        type_result = self._check_expression_impl(node, scope)
        node.inferred_type = type_result
        return type_result

    def _check_expression_impl(self, node: Expr, scope: SymbolTable) -> TypeNode:
        # print(f"DEBUG sema _check_expression_impl: {type(node).__name__}")
        # print(f"Checking expression: {type(node)} {node}")
        if isinstance(node, Literal):
            if node.type_name == '数': return self.int_type
            if node.type_name == '徽数': return self.float_type
            if node.type_name == '词': return self.string_type
            if node.type_name == '两仪': return self.bool_type
            return self.void_type
            
        elif isinstance(node, Identifier):
            symbol = scope.lookup(node.name)
            if not symbol:
                raise NameError(f"未定义的变量 '{node.name}'")
            print(f"Looking up identifier '{node.name}': {symbol.type}")
            return symbol.type
            
        elif type(node).__name__ == 'RecordLiteral':
            # Workaround for isinstance failure due to potential module reload/import issues
            # print(f"Checking RecordLiteral: {node.record_type}, types: {list(self.types.keys())}")
            # 检查记录类型是否存在
            if node.record_type not in self.types:
                raise TypeError(f"Undefined record type '{node.record_type}'")
            
            record_type = self.types[node.record_type]
            # print(f"Found record type: {record_type}")
            
            # 检查字段数量
            if len(node.field_values) != len(record_type.fields):
                raise TypeError(f"Record '{node.record_type}' expects {len(record_type.fields)} fields, got {len(node.field_values)}")
            
            # 检查每个字段的类型
            for field_name, field_expr in node.field_values.items():
                if field_name not in record_type.fields:
                    raise TypeError(f"Record '{node.record_type}' has no field '{field_name}'")
                
                expected_type = record_type.fields[field_name]
                actual_type = self.check_expression(field_expr, scope)
                
                if expected_type.name != actual_type.name:
                    raise TypeError(f"Field '{field_name}' expected type {expected_type.name}, got {actual_type.name}")
            
            return record_type

        elif isinstance(node, FieldAccess):
            # 检查记录表达式的类型
            record_type = self.check_expression(node.record, scope)
            
            if not isinstance(record_type, RecordType):
                raise TypeError(f"Cannot access field '{node.field_name}' on non-record type {record_type.name}")
            
            if node.field_name not in record_type.fields:
                raise TypeError(f"Record '{record_type.name}' has no field '{node.field_name}'")
            
            # Store the record type name for codegen (as a workaround for inferred_type issues)
            node.record_type_name = record_type.name
            # print(f"DEBUG sema: Set record_type_name to '{record_type.name}' for FieldAccess(field_name='{node.field_name}')")
            
            return record_type.fields[node.field_name]
            
        elif isinstance(node, IfExpr):
            cond_type = self.check_expression(node.condition, scope)
            if cond_type.name != '两仪':
                raise TypeError(f"条件必须是两仪，但得到 {cond_type.name}")
            
            then_type = self.check_block(node.then_branch, scope)
            if node.else_branch:
                if isinstance(node.else_branch, Block):
                    else_type = self.check_block(node.else_branch, scope)
                else:
                    else_type = self.check_expression(node.else_branch, scope)
                
                if then_type.name != else_type.name:
                    raise TypeError(f"If 分支必须有相同类型：{then_type.name} vs {else_type.name}")
                return then_type
            return self.void_type # If without else 返回 Void? Or Option?

        elif isinstance(node, BinOp):
            left_type = self.check_expression(node.left, scope)
            right_type = self.check_expression(node.right, scope)
            
            if left_type.name != right_type.name:
                raise TypeError(f"二元操作符 '{node.op}' 需要相同类型的操作数，但得到 {left_type.name} 和 {right_type.name}")
            
            # 假设所有二元操作符返回相同类型（除了比较）
            # 比较操作符返回 '两仪'
            if node.op in ('==', '!=', '<', '>', '<=', '>='):
                return self.bool_type
            
            return left_type

        elif isinstance(node, ListLiteral):
            if len(node.elements) == 0:
                # 空列表 - 多态类型，现在使用 列[空]
                return TypeNode("列", [self.void_type])
            
            # 检查所有元素类型相同
            elem_types = [self.check_expression(elem, scope) for elem in node.elements]
            first_type = elem_types[0]
            for et in elem_types[1:]:
                if et.name != first_type.name:
                    raise TypeError(f"列表元素必须有相同类型：{first_type.name} vs {et.name}")
            
            return TypeNode("列", [first_type])

        elif isinstance(node, TupleLiteral):
            elem_types = [self.check_expression(elem, scope) for elem in node.elements]
            # 元组类型是 元[Type1, Type2, ...]
            return TypeNode("元", elem_types)

        elif isinstance(node, CallExpr):
            # 检查函数调用
            func_type = self.check_expression(node.func, scope)
            
            if not isinstance(func_type, FunctionType):
                 # Maybe it's a generic type or something else?
                 # For now, assume simple function calls
                 # raise TypeError(f"Expression is not a function: {func_type.name}")
                 # Allow calling anything for now if type is unknown/void to avoid strictness issues during MVP
                 pass
            
            if isinstance(func_type, FunctionType):
                if len(node.args) != len(func_type.param_types):
                    raise TypeError(f"Function expects {len(func_type.param_types)} arguments, got {len(node.args)}")
                
                # print(f"DEBUG sema CallExpr: Checking {len(node.args)} arguments")
                for i, arg in enumerate(node.args):
                    # print(f"DEBUG sema CallExpr: Arg {i}: {type(arg).__name__} - {arg}")
                    arg_type = self.check_expression(arg, scope)
                    expected_type = func_type.param_types[i]
                    
                    # Simple type check
                    if arg_type.name != expected_type.name and expected_type.name != "空": # Void as 'Any'
                        raise TypeError(f"Argument {i+1} type mismatch: expected {expected_type.name}, got {arg_type.name}")
                
                return func_type.return_type
            
            return self.void_type

        return self.void_type

    def check_block(self, node: Block, scope: SymbolTable) -> TypeNode:
        # 块返回最后一个语句的类型
        if not node.statements:
            return self.void_type
        
        # print(f"DEBUG sema check_block: Processing {len(node.statements)} statements")
        last_type = self.void_type
        for i, stmt in enumerate(node.statements):
            # print(f"DEBUG sema check_block: Statement {i}: {type(stmt).__name__}")
            if isinstance(stmt, Expr):
                last_type = self.check_expression(stmt, scope)
            elif isinstance(stmt, VariableDef):
                # 局部变量定义
                self.check_definition(stmt, scope)
                last_type = self.void_type # 定义语句返回 Void
            else:
                # 其他语句
                pass
                
        return last_type
